import numpy as np
import cv2
# import tensorflow as tf
# from tensorflow.keras import layers, models
# from sklearn.model_selection import train_test_split
import os

# 数据生成函数（建议在单独的数据准备阶段运行）
def generate_digits_dataset(save_path="digits_dataset", img_size=28):
    # 创建保存目录
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    
    # 七段管各段坐标定义（基于28x28图像）
    segments = [
        [(7, 3), (21, 6)],     # 上
        [(3, 7), (6, 19)],     # 左上
        [(22, 7), (25, 19)],   # 右上 
        [(7, 20), (21, 23)],   # 中
        [(3, 24), (6, 36)],    # 左下
        [(22, 24), (25, 36)],  # 右下
        [(7, 37), (21, 40)]    # 下
    ]
    
    # 数字段激活配置
    digit_segments = {
        0: [1,1,1,0,1,1,1],
        1: [0,0,1,0,0,1,0],
        2: [1,0,1,1,1,0,1],
        3: [1,0,1,1,0,1,1],
        4: [0,1,1,1,0,1,0],
        5: [1,1,0,1,0,1,1],
        6: [1,1,0,1,1,1,1],
        7: [1,0,1,0,0,1,0],
        8: [1,1,1,1,1,1,1],
        9: [1,1,1,1,0,1,1]
    }
    
    # 生成500个样本/数字
    for digit in digit_segments:
        for _ in range(500):
            img = np.zeros((img_size, img_size), dtype=np.uint8)
            
            # 添加随机形变和噪声
            scale_x = np.random.uniform(0.8, 1.2)
            scale_y = np.random.uniform(0.8, 1.2)
            rotation = np.random.uniform(-15, 15)
            
            # 绘制激活段
            for i, is_on in enumerate(digit_segments[digit]):
                if is_on:
                    (x1, y1), (x2, y2) = segments[i]
                    # 应用仿射变换
                    pts = np.array([[x1, y1], [x2, y1], [x2, y2], [x1, y2]], dtype=np.float32)
                    center = np.mean(pts, axis=0)
                    
                    # 随机变换矩阵
                    M = cv2.getRotationMatrix2D(tuple(center), rotation, 1)
                    M[:, 0] *= scale_x
                    M[:, 1] *= scale_y
                    
                    pts_t = cv2.transform(pts.reshape(1, -1, 2), M).astype(int)
                    cv2.fillPoly(img, [pts_t], 255)
            
            # 添加噪声
            noise = np.random.normal(0, 25, (img_size, img_size))
            img = np.clip(img.astype(np.float32) + noise, 0, 255).astype(np.uint8)
            
            # 保存图像
            cv2.imwrite(f"{save_path}/{digit}_{_}.png", img)

# 定义CNN模型
def create_cnn_model(input_shape=(28, 28, 1)):
    model = models.Sequential([
        layers.Conv2D(16, (3,3), activation='relu', input_shape=input_shape),
        layers.MaxPooling2D((2,2)),
        layers.Conv2D(32, (3,3), activation='relu'),
        layers.MaxPooling2D((2,2)),
        layers.Flatten(),
        layers.Dense(64, activation='relu'),
        layers.Dropout(0.5),
        layers.Dense(10, activation='softmax')
    ])
    
    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])
    return model

# 训练模型（需要先运行generate_digits_dataset生成数据）
def train_model(dataset_path="digits_dataset"):
    # 加载数据集
    X = []
    y = []
    for filename in os.listdir(dataset_path):
        if filename.endswith(".png"):
            img = cv2.imread(os.path.join(dataset_path, filename), cv2.IMREAD_GRAYSCALE)
            img = cv2.resize(img, (28, 28))
            X.append(img.reshape(28, 28, 1))
            y.append(int(filename[0]))
    
    X = np.array(X, dtype=np.float32) / 255.0
    y = np.array(y)
    
    # 划分训练测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
    
    # 创建模型
    model = create_cnn_model()
    model.fit(X_train, y_train, epochs=15, validation_data=(X_test, y_test))
    model.save("digit_classifier.h5")

# 使用训练好的模型进行识别
def recognize_digits_with_cnn(img_path):
    # 加载模型
    model = tf.keras.models.load_model("digit_classifier.h5")
    
    # 原图处理流程（保持与问题中的预处理一致）
    img = cv2.imread(img_path)
    new_size = (600, 400)
    img = cv2.resize(img, new_size)
    hsv_img = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    
    # ... [保持原有预处理步骤直到获得selected_contours] ...
    # 注意：此处需要保持与原始问题代码相同的预处理步骤，此处省略重复代码
    
    # 对每个检测到的数字区域进行分类
    digits = []
    for (x, y, w, h) in selected_contours:
        roi = thresh[y:y+h, x:x+w]
        
        # 预处理ROI适配模型输入
        roi = cv2.resize(roi, (28, 28))
        roi = cv2.erode(roi, np.ones((2,2)))
        roi = cv2.dilate(roi, np.ones((2,2)))
        roi = roi.astype(np.float32)/255.0
        roi = np.expand_dims(roi, axis=(0, -1))
        
        # 预测
        pred = model.predict(roi)
        digit = np.argmax(pred)
        digits.append(digit)
        
        # 标注结果
        cv2.rectangle(img, (x, y), (x+w, y+h), (0, 255, 0), 1)
        cv2.putText(img, str(digit), (x-10, y-10), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.65, (0, 255, 0), 2)
    
    cv2.imshow("CNN Recognition", img)
    cv2.waitKey(0)

# 使用流程：
# 1. 生成数据集（首次需要运行）
generate_digits_dataset()
# 2. 训练模型
# train_model()
# 3. 实际识别
# recognize_digits_with_cnn('IMG_1813.JPG')